package com.datastructure.test.implementofdictionarytree;

public class DictionaryTree {

    public static void main(String[] args) {

    }

    public String[] trieU (String[][] operators) {
        //先计算结果集长度，并进行初始化
        int len=0;
        for(String[] opera:operators){
            if(opera[0].equals("3")||opera[0].equals("4")){
                len++;
            }
        }
        String[] res=new String[len];
        Trie trie=new Trie();
        int id=0;

        for(String[] opera:operators){
            if(opera[0].equals("1")){
                //添加单词
                trie.insert(opera[1]);
            }
            else if(opera[0].equals("2")){
                //删除单词
                trie.delete(opera[1]);
            }
            else if(opera[0].equals("3")){
                //查询单词是否存在
                res[id++]=trie.search(opera[1])?"YES":"NO";
            }
            else if(opera[0].equals("4")){
                //查找以word为前缀的单词数量
                String preNumber=String.valueOf(trie.prefixNumber(opera[1]));
                res[id++]=preNumber;
            }
        }
        return res;
    }

    //数组结构字典树
    class Trie{

        //二维数组总行数
        private final static int N=700000;
        //每一行相当于一个节点，行与行之间存在映射关系
        int[][] trie;
        //当前行插入次数
        int[] cnt;
        //当前行被标记为结尾的次数
        int[] end;
        //当前行数
        int index;

        Trie(){
            trie=new int[N][26];
            cnt=new int[N];
            end=new int[N];
            index=0;
        }


        //添加单词
        void insert(String word){
            int cur=0;
            char[] arr=word.toCharArray();
            for(char c:arr){
                //如果子节点不存在，则新建
                if(trie[cur][c-'a']==0){
                    trie[cur][c-'a']=++index;
                }
                //往子节点方向移动
                cur=trie[cur][c-'a'];
                cnt[cur]++;
            }
            end[cur]++;
        }

        void delete(String word){
            int cur=0;
            char[] arr=word.toCharArray();
            for(char c:arr){
                //往子节点方向移动，将访问次数减一
                cur=trie[cur][c-'a'];
                cnt[cur]--;
            }
            end[cur]--;

        }

        boolean search(String word){
            int cur=0;
            char[] arr=word.toCharArray();
            for(char c:arr){
                //如果子节点不存在，说明不存在该单词
                if(trie[cur][c-'a']==0){
                    return false;
                }
                cur=trie[cur][c-'a'];
            }

            //如果前面的节点都存在，并且当前行之前被标记为结尾，则存在该单词
            return end[cur]>0;
        }

        int prefixNumber(String pre){
            int cur=0;
            char[] arr=pre.toCharArray();
            for(char c:arr){
                //如果子节点不存在，说明不存在该前缀
                if(trie[cur][c-'a']==0){
                    return 0;
                }
                cur=trie[cur][c-'a'];
            }

            //返回以该单词为前缀的数量
            return cnt[cur];
        }
    }
}
